/**
 * 对内存页的分配与释放，伙伴系统主要的功能实现。
 */
#include <cds/list.h>
#include <mm/pages.h>
#include <mm/compound_pages.h>
#include <mm/page_alloc.h>
#include <mm/page_flags.h>
#include <mm/early_alloc.h>
#include <mm/memory.h>


/**
 * @brief 判断一个页是否为另一个页的伙伴。
 * 
 * 满足以下条件的会被判断为伙伴页：
 * -    两个页可以被融合成更大的复合页。
 * -    两个页的 order 与给定的 order 相同。
 * -    两个页都可以被加入伙伴系统的空闲链表中。
 * 
 * @param page 当前页的页结构体。
 * @param buddy 伙伴页的页结构体。
 * @param order 伙伴页大小的阶数。
 * 
 * @retval true 两者互为伙伴页。
 * @retval false 两者不为伙伴页。
 */
static int __is_page_buddy(
    struct page_struct *page,
    struct page_struct *buddy,
    unsigned long order
)
{
    if (!are_pages_mergeable(page, buddy)) {
        return 0;
    }
    if (page->order != order) {
        return 0;
    }
    if (!PageBuddy(page) || !PageBuddy(buddy)) {
        return 0;
    }
    return 1;
}

/**
 * @brief 对一个属于伙伴系统的复合页进行分割。
 * 
 * 进行分割的复合页必须有效且按照其阶数对齐，并且进行分割的页阶数必须大于0。
 * 
 * 操作前页需要先被移除空闲链表。
 * 
 * @param page 进行分割的复合页。
 * 
 * @return struct page_struct* 分割成功返回分割后的伙伴页页结构体，失败返回NULL。
 */
static struct page_struct *__split_page_buddy(struct page_struct *page)
{
    if (!is_page_valid(page) || !PageBuddy(page)) {
        return NULL;
    }
    if (!list_empty(&page->lru)) {
        return NULL;
    }
    struct page_struct *buddy = split_compound_page(page);
    if (buddy == NULL) {
        return NULL;
    }
    SetPageBuddy(buddy);
    return buddy;
}

/**
 * @brief 将一对伙伴页融合成一个更大的复合页。
 * 
 * 给定的两个页必须被判断为互为伙伴页。
 * 
 * 操作前页需要先被移除空闲链表。
 * 
 * @param page 当前页的页结构体。
 * @param buddy 伙伴页的页结构体。
 * @param order 伙伴页大小的阶数。
 * 
 * @return struct page_struct* 融合成功返回融合后的符合页页结构体，失败返回NULL。
 */
static struct page_struct *__merge_page_buddy(
    struct page_struct *page,
    struct page_struct *buddy,
    unsigned long order
)
{
    if (!__is_page_buddy(page, buddy, order)) {
        return NULL;
    }
    if (!list_empty(&page->lru) || !list_empty(&buddy->lru)) {
        return NULL;
    }
    return merge_pages(page, buddy);
}

/**
 * @brief 将一个特定的页从伙伴系统的空闲链表中取出。
 * 
 * 将一个 order 大小的伙伴页从其所在空闲链表中移除，并修改相关的计数。必须确保输入参数有效，
 * 否则可能发生异常。
 * 
 * @param page 被取出的页结构体。
 * @param order 被取出页大小的阶数。
 */
static void __pull_page_from_freelist(
    struct page_struct *page,
    unsigned long order
)
{
    struct free_area *free_area = page->zone->free_area + order;
    
    list_del_init(&page->lru);
    atomic_fetch_sub(&free_area->free_count, 1);
}

/**
 * @brief 将一个特定的页放入伙伴系统的空闲链表中。
 * 
 * 将一个 order 大小的伙伴页加入其对应的空闲链表，并修改相关的计数。必须确保输入参数有效，
 * 否则可能发生异常。
 * 
 * @param page 被放入的页结构体。
 * @param order 被放入页大小的阶数。
 */
static void __put_page_to_freelist(
    struct page_struct *page,
    unsigned long order
)
{
    struct free_area *free_area = page->zone->free_area + order;
    
    list_add_tail(&free_area->free_list, &page->lru);
    atomic_fetch_add(&free_area->free_count, 1);
}

/**
 * @brief 从伙伴系统的空闲链表获得一个伙伴页。
 * 
 * 该函数不会对伙伴页进行分割操作。只有在对应 order 的空闲链表不为空时才能成功分配。
 * 
 * 必须确保输入参数有效，否则可能发生异常。
 * 
 * @param zone 空闲链表所在的内存区域结构体。
 * @param order 需要获得的伙伴页大小的阶数。
 * 
 * @return struct page_struct* 分配成功返回分配到的伙伴页地址，失败返回NULL。
 */
static struct page_struct *__get_page_from_freelist(
    struct zone_struct *zone,
    unsigned long order
)
{
    if (zone == NULL || order >= PAGE_MAX_ORDER) {
        return NULL;
    }
    struct free_area *free_area = zone->free_area + order;
    if (atomic_load(&free_area->free_count) == 0) {
        return NULL;
    }
    
    struct page_struct *page = list_first_entry_or_null(
        &free_area->free_list, 
        struct page_struct, 
        lru
    );

    /// @todo 如果这里获得空闲页失败，说明出现了bug。
    if (page == NULL) {
while (1);
        return NULL;
    }
    __pull_page_from_freelist(page, order);
    return page;
}

/**
 * @brief 从伙伴系统的空闲链表中，通过分割更大的复合页获得一个伙伴页。
 * 
 * 通过分割伙伴系统中更大的复合页得到某个 order 的复合页。如果order过大，或给定内存区域的空闲
 * 链表中没有更大的页，获取失败。
 * 
 * 必须确保输入参数有效，否则可能发生异常。
 * 
 * @param zone 空闲链表所在的内存区域结构体。
 * @param order 需要获得的伙伴页大小的阶数。
 * 
 * @return struct page_struct* 分配成功返回分配到的伙伴页地址，失败返回NULL。
 */
static struct page_struct *__get_page_by_split(
    struct zone_struct *zone,
    unsigned long order
)
{
    unsigned long tmp_order = order + 1;
    struct page_struct *page = NULL;

    /// 尝试获得更大的组合页。
    do {
        page = __get_page_from_freelist(zone, tmp_order);
        if (page != NULL) {
            break;
        }
        tmp_order += 1;
    } while (tmp_order < PAGE_MAX_ORDER);

    if (page == NULL) {
        return NULL;
    }

    /// 进行分割。
    while (tmp_order > order) {
        struct page_struct *tail = split_compound_page(page);
        if (tail == NULL) {
while (1);
            /// @todo 一个BUG！
        }
        if (tail->order != page->order) {
while (1);
            /// @todo 一个BUG！
        }

        tmp_order -= 1;
        if (tail->order != tmp_order) {
while (1);
            /// @todo 一个BUG！
        }
        
        __put_page_to_freelist(tail, tmp_order);
    }

    return page;
}

/**
 * @brief 获得以给定内存页所对应的页为首页时，最大可能复合页的阶数。
 * 
 * @param page 页结构体。
 * 
 * @return unsigned long 为0时表示页无效或者该页不可形成复合页，大于0时为最大可能的阶数。
 */
static unsigned long __pfn_max_possible_order(struct page_struct *page)
{
    unsigned long start_pfn = get_page_pfn(page);
    if (start_pfn == 0) {
        return 0;
    }
    struct zone_struct *zone = page->zone;
    if (zone == NULL) {
        return 0;
    }
    if (start_pfn & 1) {
        return 0;
    }

    unsigned long zone_start_pfn = zone->zone_start_pfn;
    unsigned long zone_end_pfn = zone->owned_pages + zone_start_pfn;
    unsigned long order_count_mask = 1;
    unsigned long order_count = 0;
    for (unsigned long i = 1; start_pfn + i < zone_end_pfn; i++) {
        if (!PageBuddy(page + i)) {
            break;
        }
        if (order_count_mask & i) {
            order_count_mask <<= 1;
            order_count += 1;
        }
        if (order_count_mask & start_pfn) {
            break;
        }
    }
    if (order_count >= PAGE_MAX_ORDER) {
        order_count = PAGE_MAX_ORDER - 1;
    }

    return order_count;
}

int activate_buddy_system(void)
{
    /// 确保已经被分配的页不会被加入伙伴系统。
    __mark_earliest_allocated_pfn();
    
    /// 遍历所有的内存区域，尽可能将所有内存页融合成最大合法的复合页，并加入伙伴系统。
    unsigned long pos;
    struct zone_struct *zone;
    for_each_zone_current(pos, zone) {
        struct page_struct *zone_start_page =   \
            pfn_to_page(zone->zone_start_pfn);
        if (zone_start_page == NULL) {
            /// @todo 一个BUG！
            continue;
        }
        unsigned long zone_owned_pages = zone->owned_pages;
        unsigned long add_pages_count = 0;
        unsigned long current_offset = 0;
        while (current_offset < zone_owned_pages) {
            struct page_struct *current_page = zone_start_page + current_offset;
            if (!is_page_valid(current_page) || !PageBuddy(current_page)) {
                current_offset += 1;
                continue;
            }
            unsigned long order = __pfn_max_possible_order(current_page);
            if (order > 0) {
                if (create_compound_page(current_page, order) == NULL) {
                    /// @todo 一个BUG！
                    continue;
                }
            }
            __put_page_to_freelist(current_page, order);
            current_offset += 1UL << order;
        }
    }

    return 0;
}

/**
 * @brief 将整个复合页标记为属于伙伴系统。
 * 
 * 将所有属于给定页为首页的复合页的单页中伙伴系统的标志置位。
 * 
 * 该函数不会检查输入的有效性。
 * 
 * @param page 单页或者复合页的首页。
 */
static inline void __mark_pages_buddy(struct page_struct *page)
{
    unsigned long order = get_page_order(page);
    for (int i = 0; i < (1UL << order); i++) {
        SetPageBuddy(page + i);
    }
}

/**
 * @brief 将整个复合页标记为不属于伙伴系统。
 * 
 * 将所有属于给定页为首页的复合页的单页中伙伴系统的标志清除。
 * 
 * 该函数不会检查输入的有效性。
 * 
 * @param page 单页或者复合页的首页。
 */
static inline void __unmark_pages_buddy(struct page_struct *page)
{
    unsigned long order = get_page_order(page);
    for (int i = 0; i < (1UL << order); i++) {
        ClearPageBuddy(page + i);
    }
}

/// @todo for debug
static inline int __validate_whole_page(
    struct page_struct *page, 
    unsigned long order
)
{
    if (page == NULL || order >= PAGE_MAX_ORDER) {
        return 0;
    }
    if (order == 0) {
        return !PageBuddy(page) && !PageCompound(page) && !PageCompoundHead(page);
    }
    for (unsigned long i = 0; i < (1UL << order); i++) {
        if (PageBuddy(page + i)) {
            return 0;
        }
        if (!PageCompound(page + i)) {
            return 0;
        }
    }
    return PageCompoundHead(page);
}

struct page_struct *__alloc_pages(unsigned long order)
{
    if (order >= PAGE_MAX_ORDER) {
        return NULL;
    }

    unsigned long pos;
    struct zone_struct *zone;
    for_each_zone_current(pos, zone) {
        spinlock_acquire(&zone->lock);
        struct page_struct *page = __get_page_from_freelist(zone, order);
        if (page != NULL) {
            __unmark_pages_buddy(page);
            spinlock_release(&zone->lock);
            if (!__validate_whole_page(page, order)) {
while (1);
            }
            return page;
        }

        /// 尝试进行分裂操作
        page = __get_page_by_split(zone, order);
        if (page != NULL) {
            __unmark_pages_buddy(page);
            spinlock_release(&zone->lock);
            if (!__validate_whole_page(page, order)) {
while (1);
            }
            return page;
        }
        spinlock_release(&zone->lock);
    }
    return NULL;
}

/**
 * @brief 尝试进行页融合。
 * 
 * 尝试在空闲链表中寻找给定页的伙伴页并融合。
 *
 * 不会检查输入的有效性。
 * 
 * @return unsigned long 融合后的阶数。
 */
static inline unsigned long __try_pages_merging(struct page_struct *page)
{
    unsigned long order = page->order;
    struct zone_struct *zone = page->zone;
    while (order < PAGE_MAX_ORDER) {
        struct page_struct *buddy = __get_compound_buddy(page, order);
        if (!__is_page_buddy(page, buddy, order)) {
            return order;
        }
        __pull_page_from_freelist(buddy, order);
        page = merge_pages(page, buddy);
        if (page == NULL) {
            /// @todo 一个BUG！
while (1);
        }
        order += 1;
    }

    return PAGE_MAX_ORDER - 1;
}

/**
 * @brief 
 */
static inline void __merge_and_free_whole_page(
    struct page_struct *page, 
    unsigned long order
)
{
    unsigned long current_offset = 0;
    while (current_offset < (1UL << order)) {
        unsigned long tmp_order = get_page_order(page + current_offset);
        if (PageBuddy(page + current_offset)) {
            if (page[current_offset].order != tmp_order) {
                /// @todo: 一个BUG！
while (1);
            }
            __pull_page_from_freelist(page + current_offset, tmp_order);
        }
        if (tmp_order > 0) {
            unsigned long break_size =  \
                break_cmpound_page(page + current_offset);
            if (break_size != (1UL << tmp_order)) {
                /// @todo 一个BUG！
while (1);
            }
        }
        current_offset += 1UL << tmp_order;
    }

    if (order > 0) {
        if (create_compound_page(page, order) == NULL) {
            /// @todo 一个BUG！
while (1);
        }
    }

    __mark_pages_buddy(page);
    if (page->order != order) {
        /// @todo 一个BUG！
while (1);
    }

    /// 尝试进行页融合。
    unsigned long new_order = __try_pages_merging(page);
    if (new_order < order) {
        /// @todo 一个BUG！
while (1);
    }
    if (new_order > 0) {
        page = get_header_page(page);
    }
    __put_page_to_freelist(page, new_order);
}

/**
 * @brief 
 */
static inline void __split_and_free_whole_page(
    struct page_struct *page, 
    unsigned long target_order,
    unsigned long order
)
{
    if (PageBuddy(page)) {
        /// 已经存在于伙伴系统中就不用分割了。
        return;
    }
    unsigned long tmp_order = order;
    while (tmp_order > target_order) {
        struct page_struct *header_page = get_header_page(page);
        if (header_page == NULL) {
            /// @todo 一个BUG！
while (1);
        }
        if (split_compound_page(header_page) == NULL) {
            /// @todo 一个BUG！
while (1);
        }
        tmp_order -= 1;
    }

    __mark_pages_buddy(page);
    if (page->order != target_order) {
        /// @todo 一个BUG！
while (1);
    }
    __put_page_to_freelist(page, target_order);
}

void __free_pages(struct page_struct *page, unsigned long order)
{
    /// @todo 需要对阶数大于给定 order 的复合页的释放做进一步处理。
    unsigned long page_pfn = get_page_pfn(page);
    if (page_pfn == 0 || order >= PAGE_MAX_ORDER) {
        return;
    }
    unsigned long order_mask = (1UL << order) - 1;
    if (page_pfn & order_mask) {
        return;
    }
    if (page[order_mask].zone != page->zone) {
        return;
    }

    unsigned long tmp_order = get_page_order(get_header_page(page));
    if (tmp_order <= order) {
        __merge_and_free_whole_page(page, order);
    } else {
        __split_and_free_whole_page(page, order, tmp_order);
    }
}

unsigned long get_free_list_count(
    struct zone_struct *zone, 
    unsigned long order
)
{
    if (zone == NULL || order >= PAGE_MAX_ORDER) {
        return 0;
    }
    return atomic_load(&(zone->free_area[order].free_count));
}
